import json
from typing import List, Dict, Any
from app.utils.logger import Logger

logger = Logger("MemoryManager")

class MemoryManager:
    """记忆体管理类"""
    
    def __init__(self, max_length: int = 20):
        """初始化记忆体管理器"""
        self.memory: List[Dict[str, str]] = []  # 对话历史
        self.max_length = max_length  # 最大保留消息条数
    
    def add_to_memory(self, role: str, content: str) -> None:
        """将消息添加到记忆体，同时控制长度"""
        # 限制单条消息的长度
        if isinstance(content, str) and len(content) > 1000:
            content = content[:1000] + "...(内容已截断)"
        
        self.memory.append({"role": role, "content": content})
        if len(self.memory) > self.max_length:
            self.memory = self.memory[-self.max_length:]
        
        logger.info(f"添加消息到记忆体: {role}, 内容长度: {len(content)}")
    
    def get_memory(self) -> List[Dict[str, str]]:
        """获取当前记忆体"""
        return self.memory.copy()
    
    def clear_memory(self) -> None:
        """清空对话记忆"""
        self.memory = []
        logger.info("记忆体已清空")
    
    def get_memory_length(self) -> int:
        """获取记忆体中消息的数量"""
        return len(self.memory)
    
    def optimize_memory(self, max_chars: int = 3000) -> None:
        """优化记忆体，控制总长度"""
        total_memory_length = sum(len(msg.get('content', '')) for msg in self.memory)
        if total_memory_length > max_chars:
            # 清除较旧的记忆
            self.memory = self.memory[-5:]  # 只保留最近5条记忆
            logger.info(f"记忆体已优化，当前长度: {sum(len(msg.get('content', '')) for msg in self.memory)}")